import numpy as np
import gurobipy as gp
import cvxpy as cp
import time


class CVX:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = None, None, None

    def run(self, collector):
        for i in range(self.env.config.cd_episode):
            print("**************************************************************************")
            print(f"Current episode: {i}")
            start_timestamp = time.time()
            # 1. 生成卸载决策（a）
            self.choose_action()
            # 2. 进行资源分配（B, f）
            try:
                self.bandwidth, self.frequency = self.ra2(self.action)
            except cp.error.SolverError:
                continue
            # 3. 计算当前决策的总成本
            self.env.execute(action=self.action, bandwidth=self.bandwidth, frequency=self.frequency)
            cost = self.env.calc_cost()
            t_ratio, cpu_ratio = self.env.calc_utilization()
            print(f"Offloading cost: {cost}, Offloading num: {sum(self.action)}, Time Ratio:{t_ratio}, CPU Ratio:{cpu_ratio}, Time:{round(time.time() - start_timestamp, 2)}")
            # 4. 计算资源利用率
            collector.collect(cost, sum(self.action), t_ratio, cpu_ratio)
            collector.next()
        return self.action, self.bandwidth, self.frequency

    def choose_action(self):
        if self.action is None:
            self.action = np.zeros(self.env.M, dtype=int)
        else:
            cost = []
            for i in range(self.env.M):
                action = self.action.copy()
                action[i] = (0 if action[i] else 1)
                bandwidth, frequency = self.ra2(action=action)
                ue_p = self.check_timeout(action, bandwidth, frequency)
                c = self.env.calc_cost(action=action, bandwidth=bandwidth, frequency=frequency) if len(ue_p) == 0 else 9999
                cost.append(c)
            idx = np.argmin(cost)
            self.action[idx] = (0 if self.action[idx] else 1)

    def ra(self, action):
        b_opt = np.zeros(self.env.M)
        f_opt = np.zeros(self.env.M)
        for ues_idx in self.env.parts_ues:
            ues_idx_off = [i for i in ues_idx if int(action[i]) == 1]
            ues_num = len(ues_idx_off)
            cvx_model = gp.Model()
            b_part = cvx_model.addVars(ues_num, lb=0.001, ub=self.env.config.uav_bandwidth, vtype=gp.GRB.CONTINUOUS, name="b")
            f_part = cvx_model.addVars(ues_num, lb=0.001, ub=self.env.config.uav_cpu_frequency, vtype=gp.GRB.CONTINUOUS, name="f")
            x_part = cvx_model.addVars(ues_num, lb=1/self.env.config.uav_bandwidth, ub=1000, vtype=gp.GRB.CONTINUOUS, name="x")
            y_part = cvx_model.addVars(ues_num, lb=1 / self.env.config.uav_cpu_frequency, ub=1000, vtype=gp.GRB.CONTINUOUS, name="y")
            obj = []
            for i in range(ues_num):
                idx = ues_idx_off[i]
                obj.append((self.env.meta_trans_rate[idx] * x_part[i] + self.env.meta_task_cpu[idx] * y_part[i]) / self.env.meta_ue_loc[idx])
                cvx_model.addConstr(b_part[i] * x_part[i] == 1, name=f"c1#{i}")
                cvx_model.addConstr(f_part[i] * y_part[i] == 1, name=f"c2#{i}")
            cvx_model.addConstr(gp.quicksum(b_part) == self.env.config.uav_bandwidth, name="c3")
            cvx_model.addConstr(gp.quicksum(f_part) == self.env.config.uav_cpu_frequency, name="c4")
            cvx_model.setObjective(gp.quicksum(obj), sense=gp.GRB.MINIMIZE)
            cvx_model.Params.OutputFlag = 0
            cvx_model.Params.NonConvex = 2
            cvx_model.optimize()
            if cvx_model.status == gp.GRB.OPTIMAL:
                for i in range(ues_num):
                    b_opt[ues_idx[i]] = b_part[i].X
                    f_opt[ues_idx[i]] = f_part[i].X
        return b_opt, f_opt

    def ra2(self, action):
        b_opt = np.zeros(self.env.M)
        f_opt = np.zeros(self.env.M)
        for ues_idx in self.env.parts_ues:
            ues_idx_off = [i for i in ues_idx if int(action[i]) == 1]
            ues_num = len(ues_idx_off)
            if ues_num == 0:
                continue
            b_part = cp.Variable(ues_num, nonneg=True)
            f_part = cp.Variable(ues_num, nonneg=True)
            t_off = cp.multiply((self.env.meta_trans_rate[ues_idx_off] / self.env.meta_ue_loc[ues_idx_off]), cp.inv_pos(b_part))
            t_exe = cp.multiply((self.env.meta_task_cpu[ues_idx_off] / self.env.meta_ue_loc[ues_idx_off]), cp.inv_pos(f_part))
            obj = cp.Minimize(cp.sum(t_off + t_exe))
            cons = [cp.sum(b_part) == self.env.config.uav_bandwidth,
                    cp.sum(f_part) == self.env.config.uav_cpu_frequency]
            for i in range(ues_num):
                idx = ues_idx_off[i]
                cons.append(cp.multiply(self.env.meta_trans_rate[idx], cp.inv_pos(b_part)) + cp.multiply(self.env.meta_task_cpu[ues_idx_off], cp.inv_pos(f_part)) <= self.env.config.uav_hover_time)
            prob = cp.Problem(objective=obj, constraints=cons)
            prob.solve(solver=cp.SCS)
            for i, idx in enumerate(ues_idx_off):
                b_opt[idx] = b_part.value[i]
                f_opt[idx] = f_part.value[i]
        return b_opt, f_opt

    def check_timeout(self, action, bandwidth, frequency):
        with np.errstate(divide='ignore'):
            t_off = self.env.meta_trans_rate / bandwidth
            t_exe = self.env.meta_task_cpu / frequency
        for m in range(self.env.M):
            if action[m] == 0:
                t_off[m] = 0
                t_exe[m] = 0
        t_serv = t_off + t_exe
        timeout_result = (t_serv > self.env.config.uav_hover_time)
        # noinspection PyTypeChecker
        return [idx for idx, val in enumerate(timeout_result) if val]
